#!/usr/bin/env python3
"""
性能测试脚本：比较递归DP vs 迭代DP在不同STEP值下的性能

运行方法：
    python test_performance.py
"""

import time
import numpy as np
from typing import Dict
import sys
import os

# 添加当前目录到路径
sys.path.append(os.path.dirname(__file__))

def test_step_performance():
    """测试不同STEP值下的性能"""
    
    # 测试不同的STEP值
    test_steps = [0.5, 0.1, 0.05, 0.02]
    
    print("=" * 60)
    print("性能测试：迭代DP vs 不同STEP值")
    print("=" * 60)
    
    results = {}
    
    for step in test_steps:
        print(f"\n测试 STEP = {step}")
        print("-" * 30)
        
        # 修改dp_core中的STEP值
        import dp_core
        original_step = dp_core.STEP
        original_grid = dp_core.GRID
        
        # 更新STEP和GRID
        dp_core.STEP = step
        dp_core.GRID = np.round(np.arange(0.0, dp_core.HOURLY_MAX + 1e-9, step), 4)
        
        try:
            # 导入并测试
            from dp_core import _solve_one_scene
            
            # 简单的性能测试：运行一个场景
            alpha, beta = 0.3, 0.4
            k, j = 0, 0
            
            start_time = time.time()
            
            # 运行多次取平均
            n_runs = 3
            total_time = 0
            
            for _ in range(n_runs):
                run_start = time.time()
                result = _solve_one_scene(k, j, alpha, beta)
                run_time = time.time() - run_start
                total_time += run_time
            
            avg_time = total_time / n_runs
            grid_size = len(dp_core.GRID)
            
            results[step] = {
                'avg_time': avg_time,
                'grid_size': grid_size,
                'result': result
            }
            
            print(f"网格大小: {grid_size}")
            print(f"平均运行时间: {avg_time:.3f}秒")
            print(f"收益结果: {result:.2f}")
            
        except Exception as e:
            print(f"错误: {e}")
            results[step] = {'error': str(e)}
        
        finally:
            # 恢复原始值
            dp_core.STEP = original_step
            dp_core.GRID = original_grid
    
    # 性能总结
    print("\n" + "=" * 60)
    print("性能总结")
    print("=" * 60)
    
    print(f"{'STEP':<8} {'网格大小':<8} {'时间(秒)':<10} {'速度提升':<10}")
    print("-" * 40)
    
    baseline_time = None
    for step, data in results.items():
        if 'error' in data:
            print(f"{step:<8} {'错误':<8} {'-':<10} {'-':<10}")
            continue
            
        time_val = data['avg_time']
        grid_size = data['grid_size']
        
        if baseline_time is None:
            baseline_time = time_val
            speedup = "基准"
        else:
            speedup = f"{baseline_time/time_val:.1f}x" if time_val > 0 else "∞"
        
        print(f"{step:<8} {grid_size:<8} {time_val:<10.3f} {speedup:<10}")

def test_numba_compilation():
    """测试Numba编译是否正常"""
    print("\n" + "=" * 60)
    print("Numba编译测试")
    print("=" * 60)
    
    try:
        from numba import njit
        import numpy as np
        
        @njit
        def test_func(x):
            return x * 2 + 1
        
        # 第一次调用会触发编译
        start_time = time.time()
        result1 = test_func(5.0)
        compile_time = time.time() - start_time
        
        # 第二次调用使用编译后的代码
        start_time = time.time()
        result2 = test_func(5.0)
        exec_time = time.time() - start_time
        
        print(f"Numba编译成功!")
        print(f"首次调用(含编译): {compile_time:.6f}秒")
        print(f"后续调用: {exec_time:.6f}秒")
        print(f"编译加速比: {compile_time/exec_time:.1f}x")
        
        return True
        
    except ImportError:
        print("警告: Numba未安装! 请运行: pip install numba")
        return False
    except Exception as e:
        print(f"Numba测试失败: {e}")
        return False

def memory_usage_test():
    """测试内存使用情况"""
    print("\n" + "=" * 60)
    print("内存使用测试")
    print("=" * 60)
    
    try:
        import psutil
        import os
        
        # 获取当前进程
        process = psutil.Process(os.getpid())
        
        # 测试前内存
        mem_before = process.memory_info().rss / 1024 / 1024  # MB
        
        # 导入dp_core（会触发Numba编译）
        import dp_core
        
        # 测试后内存
        mem_after = process.memory_info().rss / 1024 / 1024  # MB
        
        print(f"导入前内存: {mem_before:.1f} MB")
        print(f"导入后内存: {mem_after:.1f} MB")
        print(f"内存增加: {mem_after - mem_before:.1f} MB")
        
    except ImportError:
        print("psutil未安装，跳过内存测试")
    except Exception as e:
        print(f"内存测试失败: {e}")

if __name__ == "__main__":
    print("开始性能测试...")
    
    # 1. 测试Numba编译
    numba_ok = test_numba_compilation()
    
    if numba_ok:
        # 2. 测试内存使用
        memory_usage_test()
        
        # 3. 测试不同STEP值的性能
        test_step_performance()
        
        print("\n" + "=" * 60)
        print("建议:")
        print("- 对于快速测试，使用 STEP=0.1 或 0.05")
        print("- 对于精确计算，使用 STEP=0.02")
        print("- Numba会在首次运行时编译，后续运行会更快")
        print("=" * 60)
    else:
        print("\n请先安装numba: pip install numba") 